# system_monitor_assistant.py

import logging
import psutil
from datetime import datetime
from typing import Dict, Any
from llama_index.core import Settings, PromptTemplate
from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.core.tools import FunctionTool
from llama_index.core.agent import ReActAgent

# --- 配置 ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

OLLAMA_BASE_URL = "http://localhost:11434"
LLM_MODEL = "qwen:8b"
EMBEDDING_MODEL = "nomic-embed-text:1.5"

# --- 客制化 Prompt 模板 ---
# 重点：明确指导 LLM 如何使用这些系统工具

CUSTOM_SYSTEM_PROMPT = PromptTemplate("""
你是一个名为 "SysMon" 的系统监控助手。
你的职责是准确报告本地计算机的当前状态。
你拥有以下两种能力：

1.  **get_current_time()**: 获取并返回当前的日期和时间。
2.  **get_system_usage()**: 获取并返回 CPU 和内存的使用百分比。

### 行为准则
- **直接调用**: 当用户询问时间或系统资源时，直接调用相应的工具函数。
- **精确回答**: 工具返回的数据就是答案，无需额外解释或计算。
- **格式化输出**:
  - 时间格式: YYYY年MM月DD日 HH:MM:SS
  - CPU 使用率: "CPU 使用率: XX.X%"
  - 内存使用率: "内存使用率: YY.Y%"
  - 同时询问时: 先报时间，再报 CPU，最后报内存，每项占一行。
- **禁止行为**:
  - 绝不尝试自己计算时间或系统使用率。
  - 不要编造数据。
  - 如果问题与系统状态无关，请回答：“我是一个系统监控助手，只能回答关于时间或系统资源的问题。”

### 示例
用户: 现在几点了？
助手: Thought: 用户想知道当前时间。我应该调用 get_current_time()。
       Action: get_current_time()
       Observation: 2025-08-29 08:30:15
       Answer: 2025年08月29日 08:30:15

用户: CPU 和内存怎么样？
助手: Thought: 用户想了解 CPU 和内存使用情况。我应该调用 get_system_usage()。
       Action: get_system_usage()
       Observation: {'cpu_percent': 15.3, 'memory_percent': 45.7}
       Answer: CPU 使用率: 15.3%
              内存使用率: 45.7%
""")


# --- 定义系统信息工具 (Tools) ---
def get_current_time() -> str:
    """
    获取当前的日期和时间。
    Returns:
        格式化的当前时间字符串。
    """
    now = datetime.now()
    return now.strftime("%Y-%m-%d %H:%M:%S")


def get_system_usage() -> Dict[str, float]:
    """
    获取当前的 CPU 和内存使用率。
    Returns:
        包含 'cpu_percent' 和 'memory_percent' 的字典。
    """
    cpu_percent = psutil.cpu_percent(interval=1)  # 1秒平均
    memory_percent = psutil.virtual_memory().percent
    return {
        "cpu_percent": round(cpu_percent, 1),
        "memory_percent": round(memory_percent, 1)
    }


# 将函数包装为 LlamaIndex Tool
time_tool = FunctionTool.from_defaults(
    fn=get_current_time,
    name="get_current_time",
    description="获取当前的日期和时间。当用户询问'现在几点'、'当前时间'等问题时调用此工具。"
)

system_usage_tool = FunctionTool.from_defaults(
    fn=get_system_usage,
    name="get_system_usage",
    description="获取 CPU 和内存的实时使用百分比。当用户询问'CPU 使用率'、'内存占用'、'系统负载'等问题时调用此工具。"
)

# --- 设置 LLM 和 Embedding ---
# 注意：虽然我们主要用 LLM 进行推理和工具调用，但 Settings.embed_model 在某些内部操作中可能仍被使用
Settings.llm = Ollama(
    model=LLM_MODEL,
    base_url=OLLAMA_BASE_URL,
    request_timeout=60.0,
    system_prompt=CUSTOM_SYSTEM_PROMPT.template,  # 应用客制化系统提示
)

# 使用 nomic-embed-text:v1.5 作为嵌入模型
Settings.embed_model = OllamaEmbedding(
    model_name=EMBEDDING_MODEL,
    base_url=OLLAMA_BASE_URL,
)


# --- 创建智能体 (Agent) ---
def create_agent() -> ReActAgent:
    tools = [time_tool, system_usage_tool]

    agent = ReActAgent.from_tools(
        tools=tools,
        llm=Settings.llm,
        verbose=True,  # 输出完整的推理链 (Thought/Action/Observation)
    )
    return agent


# --- 主程序 ---
def main():
    print("🖥️  SysMon 系统监控助手已启动 (qwen:8b + nomic-embed-text:1.5)")
    print("功能: 获取当前时间 | 监控 CPU/内存使用率")
    print("输入 'quit' 退出。")
    print("-" * 50)

    agent = create_agent()

    while True:
        user_query = input("\n📊 询问系统状态: ").strip()
        if user_query.lower() == "quit":
            print("SysMon 助手已关闭。")
            break
        if not user_query:
            continue

        try:
            response = agent.chat(user_query)
            # `response.response` 包含最终的 Answer 部分
            print(f"\n🤖 SysMon: {response.response}")
        except Exception as e:
            print(f"❌ 执行出错: {e}")


if __name__ == "__main__":
    main()